#pragma once

#include <memory>
#include <vector>
#include <cstdint>
#include <optional>

#include "txdnn_bindings.h"
#include "TxDNN/txdnn_graph.h"
#include "TxDNN/frontend/eager.h"
#include "TxDNN/frontend/graph_properities.h"

namespace pytxdnn {

class PyEager {
public:
  PyEager() {}

  std::shared_ptr<txdnn_frontend::graph::TensorAttributes> tensor(const std::vector<uint32_t>& dim,
                                                                  const std::vector<uint32_t>& stride,
                                                                  const txdnnDataType_t data_type,
                                                                  const bool& is_virtual, const std::string& name);

  std::shared_ptr<txdnn_frontend::graph::TensorAttributes> tensor_like(
    const std::shared_ptr<txdnn_frontend::graph::TensorAttributes>& tensor, const std::string& name);

  std::shared_ptr<txdnn_frontend::graph::TensorAttributes> tensor_like(const py::object& tensor);

  std::shared_ptr<txdnn_frontend::graph::TensorAttributes>
  matmul(std::shared_ptr<txdnn_frontend::graph::TensorAttributes>& A,
         std::shared_ptr<txdnn_frontend::graph::TensorAttributes>& B,
         const txdnnDataType_t& data_type,
         const double padding_value,
         const std::string& name);

  std::shared_ptr<txdnn_frontend::graph::TensorAttributes>
  rmsnorm(const txdnn_frontend::graph::NormFwdPhase_t forward_phase,
          std::shared_ptr<txdnn_frontend::graph::TensorAttributes>& x,
          std::shared_ptr<txdnn_frontend::graph::TensorAttributes>& scale,
          std::shared_ptr<txdnn_frontend::graph::TensorAttributes>& bias,
          std::shared_ptr<txdnn_frontend::graph::TensorAttributes>& epsilon,
          const txdnnDataType_t& data_type,
          const std::string& name);

//   bias() {
//     // return tenosr : set_name, set_output, set_data_type
//   }

  void validate();

  void build_operation_graph();

  void create_execution_plans();

  void check_support();

  void build_plans();

  void execute(std::unordered_map<int64_t, std::intptr_t> var_pack,
                std::intptr_t workspace,
                std::optional<std::intptr_t> handle);

public:
  txdnn_frontend::eager::Eager eager_;

}; // class PyEager

} // namespace pytxdnn